import torch
import torch.utils.data as tud
from torch import nn

from model_utils import AnomalyConfusionMatrix, EPOCHS, create_dataloader, device, \
    train_and_test_model


class USAD(nn.Module):
    def __init__(self, input_size, z_dim, alpha, beta):
        super(USAD, self).__init__()
        self.name = 'USAD'
        self.epoch = 1
        self.alpha = alpha
        self.beta = beta
        self.encoder = nn.Sequential(
            nn.Linear(input_size, input_size // 2),
            nn.ReLU(),
            nn.Linear(input_size // 2, input_size // 4),
            nn.ReLU(),
            nn.Linear(input_size // 4, z_dim)
        )

        self.decoder1 = nn.Sequential(
            nn.Linear(z_dim, input_size // 4),
            nn.ReLU(),
            nn.Linear(input_size // 4, input_size // 2),
            nn.ReLU(),
            nn.Linear(input_size // 2, input_size),
            nn.Sigmoid()
        )

        self.decoder2 = nn.Sequential(
            nn.Linear(z_dim, input_size // 4),
            nn.ReLU(),
            nn.Linear(input_size // 4, input_size // 2),
            nn.ReLU(),
            nn.Linear(input_size // 2, input_size),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        z = self.encoder(x)
        decoded1 = self.decoder1(z)
        decoded2 = self.decoder2(z)
        decoded12 = self.decoder2(self.encoder(decoded1))
        return (decoded1, decoded2), decoded12


def compute_usad_loss(model, x, **kwargs):
    epoch = kwargs.get('epoch', 0) + 1
    x = x.view(x.shape[0], -1)
    (decoded1, decoded2), decoded12 = model(x)
    loss1 = nn.MSELoss().to(device)(x, decoded1) / epoch + nn.MSELoss(). \
        to(device)(x, decoded12) * (1 - 1 / epoch)
    loss2 = nn.MSELoss().to(device)(x, decoded2) / epoch - nn.MSELoss(). \
        to(device)(x, decoded12) * (1 - 1 / epoch)
    # We use backward here because USAD needs to backward loss twice
    loss1.backward(retain_graph=True)
    return loss2


def get_optimizer(model):
    return torch.optim.Adam(lr=1e-4, betas=(0.9, 0.99), params=model.parameters())


def build_usad_model(input_size, **kwargs):
    model = USAD(input_size=input_size, z_dim=40,
                 alpha=0.1, beta=0.9).to(device)
    optimizer = get_optimizer(model)
    return model, optimizer


def train_and_test_usad_model(normal_dataloader: tud.DataLoader,
                              attack_dataloader: tud.DataLoader,
                              epochs: int,
                              dataset_name: str,
                              data_dir: str,
                              continue_train: bool) -> AnomalyConfusionMatrix:
    return train_and_test_model(compute_loss=compute_usad_loss,
                                build_model=build_usad_model,
                                normal_dl=normal_dataloader,
                                attack_dl=attack_dataloader,
                                epochs=epochs,
                                model_name='USAD',
                                continue_train=continue_train,
                                dataset=dataset_name,
                                data_dir=data_dir,
                                is_recur=False)


if __name__ == "__main__":
    normal_dl, attack_dl = create_dataloader(
        './data/MSL/C-1', timemode=True, reverse_label=False)

    max_confusion_matrix_with_attn_false = train_and_test_usad_model(normal_dataloader=normal_dl,
                                                                     attack_dataloader=attack_dl,
                                                                     epochs=EPOCHS,
                                                                     dataset_name='MSL/C-1')
